package com.seanLab.tool.TagBat;

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import com.seanLab.domain.Article;
import com.seanLab.domain.TagInfo;
import com.seanLab.dto.SuggestArticleKeywordsDto;
import com.seanLab.dto.SuggestModelArticleDto;
import com.seanLab.dto.SuggestModelImageDto;
import com.seanLab.tool.TagSuggestModel.ExpandRank.common.AnsjSegment;
import com.seanLab.tool.TagSuggestModel.TopicModel.MyLDAModel;
import com.seanLab.tool.TagSuggestor;
import org.apache.lucene.analysis.core.SimpleAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.StoredField;
import org.apache.lucene.document.StringField;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;

import java.io.*;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.logging.Logger;


/***
 * new project as Tagger in Data_Process
 */
@Deprecated
public class NewTagsBat {
    public static final String diretory = "workingDir";
    public static String articleDataset = "chens.article_20180306_500.json";
    public static String imageDataset = "chens.validImage_20180306_5000.json";
//    public static String indexName = "";
    public static final String resultDir = diretory + File.separator + "tagResults";
    public static final String ldaModelPath = diretory + File.separator + "LDAModel20180321" + File.separator + "MyLDAModel.json";
    public static final String vocabularyPath = diretory + File.separator + "LDAModel20180321" + File.separator + "vocabulary.dict";
    public static final String idfCorpus = diretory + File.separator + "idf.dat";
    public static int threadNum = 2;
    public static int verboseStep = 100;

    private static int tagCount = 0;
    private static final Object tagCountLock = new Object();
    private static int threadCount = 0;

    private static Logger LOG = Logger.getAnonymousLogger();

    public String findImageByImageId(IndexSearcher searcher, String id) throws IOException {
        Query query = new TermQuery(new Term("imageId", id));
        TopDocs topDocs = searcher.search(query, 1);
        if (topDocs.totalHits > 0) {
            int docId = topDocs.scoreDocs[0].doc;
            Document doc = searcher.doc(docId);
            return doc.get("json");
        }
        return null;
    }

    public IndexSearcher loadIndex() throws IOException {
        DirectoryReader directoryReader = DirectoryReader.open(
                FSDirectory.open(Paths.get(diretory + File.separator + getIndexName())));
        return new IndexSearcher(directoryReader);
    }

    private String getIndexName() {
        String indexName = imageDataset.replace(".json", "_Index");
        return indexName;
    }

    private static int[] generateTopTopics(double[] topics) {
        Integer[] indexes = new Integer[topics.length];
        for (int i = 0; i < indexes.length; i++) {
            indexes[i] = i;
        }
        Arrays.sort(indexes, new Comparator<Integer>() {
            @Override
            public int compare(final Integer i1, final Integer i2) {
                return Double.compare(topics[i2], topics[i1]);
            }
        });
        int topNum = topics.length > 10 ? 10 : topics.length;
        int[] result = new int[topNum];
        for (int i = 0; i < topNum; i++) {
            result[i] = indexes[i];
        }
        return result;
    }

    private void parralleTag(List<String> articles, MyLDAModel ldaModel) throws IOException {
        ExecutorService threadPool = Executors.newFixedThreadPool(threadNum);
        List<Callable<Object>> todo = new ArrayList<>(threadNum);
        IndexSearcher searcher = loadIndex();
        for (int i = 0; i <  threadNum; i++) {
            List<String> articleSublist = articles.subList(
                    i * articles.size() / threadNum,
                    (i + 1) * articles.size() / threadNum);
            String threadName = "Thread-" + threadCount++;
            todo.add(new Callable<Object>() {
                @Override
                public Object call() throws IOException {
                    String articlesResultPath = articleDataset.replace(".json", "_" + threadName + ".json");
                    String imagesResultPath = imageDataset.replace(".json", "_" + threadName + ".json");
                    BufferedWriter bwArticle = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(resultDir + File.separator + articlesResultPath)));
                    BufferedWriter bwImage = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(resultDir + File.separator + imagesResultPath)));

                    JsonParser jsonParser = new JsonParser();
                    Gson gson = new Gson();
                    TagSuggestor tagSuggestor = new TagSuggestor();
                    tagSuggestor.loadModel(null);

                    int count = 0;

                    for (String articleString : articleSublist) {
                        JsonObject articleJson = jsonParser.parse(articleString).getAsJsonObject();

                        List<JsonObject> imageJsonList = new ArrayList<>();
                        List<SuggestModelImageDto> imageDtoList = new ArrayList<>();

                        Article article = gson.fromJson(articleJson, Article.class);
                        for (String imgId : article.getImgID()) {
                            String imageString = findImageByImageId(searcher, imgId);
                            if (imageString != null) {
                                JsonObject image = jsonParser.parse(imageString).getAsJsonObject();
                                imageDtoList.add(new SuggestModelImageDto(image.get("offset").getAsInt(), image.get("imgDescription").getAsString(),
                                        image.get("contextAbove").getAsString(), image.get("contextBelow").getAsString()));
                                imageJsonList.add(image);
                            } else {
                                System.out.println("Can not find img of " + imgId);
                            }
                        }
                        SuggestModelArticleDto articleDto = new SuggestModelArticleDto(article.getTitle(), article.getTags(), article.getContent(), imageDtoList);

                        List<List<TagInfo>> tagsOfImages = tagSuggestor.doTag(articleDto);
                        List<SuggestArticleKeywordsDto> keywordsOfArticle = tagSuggestor.doKeyword(articleDto);

                        double[] topicDistribution = ldaModel.predictDoc(article.getContent());
                        articleJson.add("keywords", gson.toJsonTree(keywordsOfArticle));
                        articleJson.add("articleTopic", gson.toJsonTree(topicDistribution));
                        gson.toJson(articleJson, bwArticle);
                        bwArticle.newLine();
                        int[] topTopics = generateTopTopics(topicDistribution);
                        for (int i = 0; i < tagsOfImages.size(); i++) {
                            JsonObject imageJson = imageJsonList.get(i);
                            if (tagsOfImages.get(i).size() > 5) {
                                imageJson.add("imgTagList", gson.toJsonTree(tagsOfImages.get(i).subList(0, 5)));
                            }
                            else {
                                imageJson.add("imgTagList", gson.toJsonTree(tagsOfImages.get(i)));
                            }
                            imageJson.add("imgTopic", gson.toJsonTree(topicDistribution));
                            imageJson.add("topTopics", gson.toJsonTree(topTopics));
                            gson.toJson(imageJson, bwImage);
                            bwImage.newLine();
                        }

                        if (++count % verboseStep == 0) {
                            synchronized (tagCountLock) {
                                tagCount += count;
                                count = 0;
                                LOG.info("Tagged " + tagCount + " documents.(" + threadName + ")");
                            }
                        }
                    }

                    synchronized (tagCountLock) {
                        tagCount += count;
                        LOG.info("Tagged " + tagCount + " documents.(" + threadName + ")");
                    }
                    bwArticle.close();
                    bwImage.close();
                    return null;
                }
            });
        }

        try {
            threadPool.invokeAll(todo);
            threadPool.shutdown();
            LOG.info("Tagged " + tagCount + " documents.(Main)");
        } catch (InterruptedException e) {
            LOG.info("do tag failed.");
            e.printStackTrace();
        }
    }

    public void doTag() throws IOException {
        long start = System.currentTimeMillis();

        BufferedReader reader = new BufferedReader(new InputStreamReader(new FileInputStream(diretory + File.separator + articleDataset)));
        List<String> articles = new ArrayList<>(verboseStep * threadNum * 2 + 1024);
        File resultFile = new File(resultDir);
        if (!resultFile.exists()) {
            resultFile.mkdir();
        }
        LOG.info("loading lda model.");
        MyLDAModel ldaModel = new Gson().fromJson(new BufferedReader(new InputStreamReader(
                        new FileInputStream(ldaModelPath))),
                MyLDAModel.class);
        LOG.info("loading vocabulary.");
        ldaModel.loadVocabulary(vocabularyPath);
        String line;
        int count = 0;
        while ((line = reader.readLine()) != null) {
            articles.add(line);
            if (++count % verboseStep == 0) {
                LOG.info("read doc " + count);
            }
            if (count % (verboseStep * threadNum * 2) == 0) {
                parralleTag(articles, ldaModel);
                articles.clear();
            }
        }
        reader.close();
        if (!articles.isEmpty()) {
            parralleTag(articles, ldaModel);
        }

        long end = System.currentTimeMillis();
        LOG.info("total doc: " + count);
        LOG.info("Do Tag: " + String.valueOf(end - start) + "ms with " + threadNum + " thread");
    }

    public void buildImageIndex() throws IOException{
        long start = System.currentTimeMillis();
        String modelPath = diretory + File.separator + getIndexName();
        File modelDir = new File(modelPath);
        if (!modelDir.exists()) {
            modelDir.mkdir();
        } else {
            modelDir.delete();
            modelDir.mkdir();
        }
        Directory indexDir = FSDirectory.open(Paths.get(modelDir.getPath()));
        IndexWriter indexWriter = new IndexWriter(indexDir, new IndexWriterConfig(new SimpleAnalyzer()));

        JsonParser jsonParser = new JsonParser();
        BufferedReader reader = new BufferedReader(new InputStreamReader(new FileInputStream(diretory + File.separator + imageDataset)));

        String line;
        int count = 0;
        while ((line = reader.readLine()) != null) {
            JsonObject object = jsonParser.parse(line).getAsJsonObject();
            String imageId = object.get("_id").getAsString();
            String articleId = object.get("articleID").getAsString();

            Document doc = new Document();
            doc.add(new StringField("imageId", imageId, Field.Store.NO));
            doc.add(new StringField("articleId", articleId, Field.Store.NO));
            doc.add(new StoredField("json", line));
            indexWriter.addDocument(doc);

            if (++count % verboseStep == 0) {
                LOG.info("index image " + count);
            }
        }
        reader.close();
        indexWriter.close();
        LOG.info("total index image " + count);
        long end = System.currentTimeMillis();
        System.out.println("Build Index: " + String.valueOf(end - start) + "ms");
    }

    private void parralleIDF(List<String> articles, ConcurrentHashMap<String, Integer> idf) throws IOException {
        ExecutorService threadPool = Executors.newFixedThreadPool(threadNum);
        List<Callable<Object>> todo = new ArrayList<>(threadNum);

        for (int i = 0; i <  threadNum; i++) {
            List<String> articleSublist = articles.subList(
                    i * articles.size() / threadNum,
                    (i + 1) * articles.size() / threadNum);
            String threadName = "Thread-" + threadCount++;
            todo.add(new Callable<Object>() {
                @Override
                public Object call() throws IOException {
                    JsonParser jsonParser = new JsonParser();
                    AnsjSegment segment = new AnsjSegment();
                    int count = 0;

                    for (String articleString : articleSublist) {
                        JsonObject articleJson = jsonParser.parse(articleString).getAsJsonObject();
                        String text = "" + articleJson.get("title").getAsString() + articleJson.get("content").getAsString();
                        HashSet<String> wordSet = new HashSet<>();
                        wordSet.addAll(Arrays.asList(segment.segment(text)));
                        for (String word : wordSet) {
                            idf.put(word, idf.getOrDefault(word, 0) + 1);
                        }

                        if (++count % verboseStep == 0) {
                            synchronized (tagCountLock) {
                                tagCount += count;
                                count = 0;
                                LOG.info("IDF " + tagCount + " documents.(" + threadName + ")");
                            }
                        }
                    }

                    synchronized (tagCountLock) {
                        tagCount += count;
                        LOG.info("IDF " + tagCount + " documents.(" + threadName + ")");
                    }
                    return null;
                }
            });
        }

        try {
            threadPool.invokeAll(todo);
            threadPool.shutdown();
            LOG.info("IDF " + tagCount + " documents.(Main)");
        } catch (InterruptedException e) {
            LOG.info("do IDF failed.");
            e.printStackTrace();
        }
    }

    public void buildIDF() throws IOException{
        long start = System.currentTimeMillis();

        BufferedReader reader = new BufferedReader(new InputStreamReader(new FileInputStream(diretory + File.separator + articleDataset)));
        List<String> articles = new ArrayList<>(verboseStep * threadNum * 2 + 1024);
        ConcurrentHashMap<String, Integer> idf = new ConcurrentHashMap<>();
        String line;
        int count = 0;
        while ((line = reader.readLine()) != null) {
            articles.add(line);
            if (++count % verboseStep == 0) {
                LOG.info("read doc " + count);
            }
            if (count % (verboseStep * threadNum * 2) == 0) {
                parralleIDF(articles, idf);
                articles.clear();
            }
        }
        reader.close();
        if (!articles.isEmpty()) {
            parralleIDF(articles, idf);
        }

        BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(idfCorpus)));
        writer.write(count + "\n");
        for (Map.Entry<String, Integer> item : idf.entrySet()) {
            writer.write(item.getKey() + "\t" + item.getValue() + "\n");
        }
        writer.close();

        long end = System.currentTimeMillis();
        LOG.info("total doc: " + count);
        LOG.info("Build Dict: " + String.valueOf(end - start) + "ms with " + threadNum + " thread");

    }

    public static void main(String[] args) throws IOException{
        System.out.println("Good Luck!");
        String usage = "Usage: java -jar xx.jar $threadNum $articleDataset $imageDataset $action[index|tag|index-tag|idf] $verboseStep";
        if (args.length != 5) {
            System.out.println(usage);
            return;
        }
        threadNum = Integer.parseInt(args[0]);
        articleDataset = args[1];
        imageDataset = args[2];
        System.out.println(threadNum + " " + articleDataset + " " + imageDataset);
        String action = args[3];
        verboseStep = Integer.parseInt(args[4]);
        if (action.equals("index")) {
            new NewTagsBat().buildImageIndex();
        }
        else if (action.equals("tag")) {
            new NewTagsBat().doTag();
        }
        else if (action.equals("index-tag")) {
            NewTagsBat app = new NewTagsBat();
            app.buildImageIndex();
            app.doTag();
        }
        else if (action.equals("idf")) {
            new NewTagsBat().buildIDF();
        }
        else {
            System.out.println(usage);
        }
    }

    private static class BuildImageIndex {
        public static void main(String[] args) throws IOException {
            new NewTagsBat().buildImageIndex();
        }
    }

    private static class TestIndex {
        public static void main(String[] args) throws IOException {
            String testId = "imageID-c833cd1e-a50e-4476-bbd3-73b439471259";
            NewTagsBat app = new NewTagsBat();
            IndexSearcher searcher = app.loadIndex();
            String json = app.findImageByImageId(searcher, testId);
            System.out.println(json);
        }
    }

    private static class TestTag {
        public static void main(String[] args) throws IOException {
            new NewTagsBat().doTag();
        }
    }

    private static class TestIDF {
        public static void main(String[] args) throws IOException {
            new NewTagsBat().buildIDF();
        }
    }

    private static class TestTopTopics {
        public static void main(String[] args) {
            double[] a1 = new double[] {0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
            double[] a2 = new double[] {0.0, 0.0, 0.3, 1.0, 0.0, 0.0, 0.0, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
            double[] a3 = new double[] {0.0, 0.0, 0.3, 1.0, 0.0, 0.0, 0.0, 0.0};
            System.out.println(Arrays.toString(generateTopTopics(a1)));
            System.out.println(Arrays.toString(generateTopTopics(a2)));
            System.out.println(Arrays.toString(generateTopTopics(a3)));
        }
    }
}
